package org.example.distributed_lock;

import org.apache.zookeeper.*;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

/**
 * 使用 zookeeper 模拟实现分布式锁
 */
public class DistributedLock {
    private String connectString = "47.108.48.236:2182";
    private int sessionTimeout = 300000;
    private ZooKeeper zooKeeper;
    private String path = "/lock";
    private CountDownLatch contentLatch = new CountDownLatch(1);
    private CountDownLatch lockLatch = new CountDownLatch(1);

    private String waitPath;
    private String curNode;
    /**
     * 与 zookeeper 建立连接, 若 /lock 不存在, 则创建根节点
     */
    public DistributedLock() {
        try {
            zooKeeper = new ZooKeeper(connectString, sessionTimeout, new Watcher() {
                @Override
                public void process(WatchedEvent watchedEvent) {
                    // 连接成功
                    if (watchedEvent.getState() == Event.KeeperState.SyncConnected) {
                        contentLatch.countDown();
                    }
                    // 前一个节点释放锁
                    if (watchedEvent.getType() == Event.EventType.NodeDeleted && watchedEvent.getPath().equals(waitPath)) {
                        System.out.println(curNode + " 获取到锁...");
                        lockLatch.countDown();
                    }
                }
            });
            contentLatch.await();
            System.out.println("zookeeper 连接成功...");
            if (zooKeeper.exists(path, null) == null) {
                zooKeeper.create(path, path.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } catch (KeeperException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 进行加锁
     * 监听前一个节点(需要考虑异常情况...)
     */
    public void zkLock() throws InterruptedException, KeeperException {
        // 创建临时顺序节点
        curNode = zooKeeper.create(path + "/" + "seq-", null, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
        // 获取当前 /lock 下的节点
        List<String> keeperChildren = zooKeeper.getChildren(path, null);
        // 若只有一个节点, 则当前可以加锁
        if (keeperChildren.size() == 1) {
            System.out.println(curNode + " 获取到锁...");
        } else {
            // 计算当前节点位置
            Collections.sort(keeperChildren);
            String thisNode = curNode.substring((path + "/").length());
            int curIndex = keeperChildren.indexOf(thisNode);
            if (curIndex > 0){
                // 监听前一个节点
                waitPath = path + "/" + keeperChildren.get(curIndex - 1);
                zooKeeper.getData(waitPath, true, null);
                // 等待前一个节点释放锁
                lockLatch.await();
            }else if (curIndex == 0) {
                System.out.println(curNode + " 获取到锁...");
            }else {
                throw new RuntimeException("下标获取失败: node: " + thisNode + " children: " + keeperChildren);
            }
        }
    }

    /**
     * 释放锁
     */
    public void zkUnlock() throws InterruptedException, KeeperException {
        System.out.println(curNode + " 释放锁...");
        zooKeeper.delete(curNode, -1);
    }
}
